import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

from src.data.loader import load_data, prepare_time_series_data
from src.visualization.plots import (
    plot_energy_trends, plot_energy_distribution, plot_hourly_consumption_production,
    plot_stacked_area, plot_lag_features, plot_correlation_matrix, plot_model_predictions
)
from src.models.rnn_model import create_rnn_model, train_rnn_model, predict_rnn
from src.models.lstm_model import create_lstm_model, train_lstm_model, predict_lstm
from src.utils.helpers import evaluate_model, print_evaluation_metrics

def main():
    # 设置随机种子以确保结果可重现
    np.random.seed(42)
    
    # 创建必要的目录
    os.makedirs('e:/Desktop/利用RNN和LSTM预测小时用电和发电量/总实验/models', exist_ok=True)
    os.makedirs('e:/Desktop/利用RNN和LSTM预测小时用电和发电量/总实验/data', exist_ok=True)
    
    # 数据文件路径 - 检查data目录中是否有文件，如果没有则使用原始路径
    data_file_in_data = 'e:/Desktop/利用RNN和LSTM预测小时用电和发电量/总实验/data/electricityConsumptionAndProductioction.csv'
    original_data_file = 'e:/Desktop/利用RNN和LSTM预测小时用电和发电量/总实验/electricityConsumptionAndProductioction.csv'
    
    if os.path.exists(data_file_in_data):
        data_file = data_file_in_data
    elif os.path.exists(original_data_file):
        data_file = original_data_file
        # 复制文件到data目录
        import shutil
        os.makedirs(os.path.dirname(data_file_in_data), exist_ok=True)
        shutil.copy2(original_data_file, data_file_in_data)
        print(f"已将数据文件复制到 {data_file_in_data}")
        data_file = data_file_in_data
    else:
        # 尝试在其他位置查找
        possible_locations = [
            '/home/dataset/electricityConsumptionAndProductioction.csv',
            'electricityConsumptionAndProductioction.csv'
        ]
        for loc in possible_locations:
            if os.path.exists(loc):
                data_file = loc
                print(f"找到数据文件: {data_file}")
                break
        else:
            raise FileNotFoundError("找不到数据文件，请确保数据文件存在于正确的位置")
    
    # 加载数据
    print("加载数据...")
    df = load_data(data_file)
    print(f"数据加载完成，共 {len(df)} 条记录")
    
    # 定义能源列表
    energy_sources = ['Nuclear', 'Wind', 'Hydroelectric', 'Oil and Gas', 'Coal', 'Solar', 'Biomass']
    
    # 数据探索
    print("\n数据探索...")
    
    # 1. 探索电力的发展趋势
    print("绘制能源趋势图...")
    plot_energy_trends(df, energy_sources)
    
    # 2. 发电的分布
    print("绘制能源分布图...")
    plot_energy_distribution(df, energy_sources)
    
    # 3. 每小时用电量和发电量的线图
    print("绘制每小时用电量和发电量线图...")
    plot_hourly_consumption_production(df)
    
    # 4. 显示每小时的电力生产的来源与堆叠面积图
    print("绘制堆叠面积图...")
    plot_stacked_area(df, energy_sources)
    
    # 5. 滞后图
    print("绘制滞后图...")
    plot_lag_features(df, 'Consumption')
    
    # 6. 电力生产与消费数据集相关矩阵
    print("绘制相关矩阵...")
    plot_correlation_matrix(df, ['Consumption', 'Production'] + energy_sources)
    
    # 准备时间序列数据
    print("\n准备时间序列数据...")
    sequence_length = 24  # 使用前24小时预测下一小时
    train_generator, test_generator, scaler, train_data, test_data = prepare_time_series_data(
        df, 'Consumption', sequence_length=sequence_length
    )
    
    # RNN模型训练和预测
    print("\n训练RNN模型...")
    rnn_model = create_rnn_model(sequence_length)
    rnn_history, rnn_model = train_rnn_model(
        rnn_model, train_generator, test_generator, 
        epochs=50, 
        model_path='e:/Desktop/利用RNN和LSTM预测小时用电和发电量/总实验/models/rnn_model.h5'
    )
    
    print("使用RNN模型进行预测...")
    rnn_predictions = predict_rnn(rnn_model, test_generator, test_data, sequence_length)
    
    # 评估RNN模型
    rnn_metrics = evaluate_model(test_data, rnn_predictions)
    print("\nRNN模型评估:")
    print_evaluation_metrics(rnn_metrics)
    
    # 绘制RNN预测结果
    print("绘制RNN预测结果...")
    plot_model_predictions(test_data, rnn_predictions, 'RNN模型: 实际值 vs 预测值 (用电量)', scaler)
    
    # LSTM模型训练和预测
    print("\n训练LSTM模型...")
    lstm_model = create_lstm_model(sequence_length)
    lstm_history, lstm_model = train_lstm_model(
        lstm_model, train_generator, test_generator, 
        epochs=50, 
        model_path='e:/Desktop/利用RNN和LSTM预测小时用电和发电量/总实验/models/lstm_model.h5'
    )
    
    print("使用LSTM模型进行预测...")
    lstm_predictions = predict_lstm(lstm_model, test_generator, test_data, sequence_length)
    
    # 评估LSTM模型
    lstm_metrics = evaluate_model(test_data, lstm_predictions)
    print("\nLSTM模型评估:")
    print_evaluation_metrics(lstm_metrics)
    
    # 绘制LSTM预测结果
    print("绘制LSTM预测结果...")
    plot_model_predictions(test_data, lstm_predictions, 'LSTM模型: 实际值 vs 预测值 (用电量)', scaler)
    
    # 比较RNN和LSTM模型
    print("\n比较RNN和LSTM模型:")
    metrics_comparison = pd.DataFrame({
        'RNN': [rnn_metrics['MSE'], rnn_metrics['RMSE'], rnn_metrics['MAE'], rnn_metrics['R2']],
        'LSTM': [lstm_metrics['MSE'], lstm_metrics['RMSE'], lstm_metrics['MAE'], lstm_metrics['R2']]
    }, index=['MSE', 'RMSE', 'MAE', 'R2'])
    
    print(metrics_comparison)
    
    # 绘制比较图
    plt.figure(figsize=(12, 6))
    metrics_comparison.plot(kind='bar')
    plt.title('RNN vs LSTM 模型性能比较')
    plt.ylabel('指标值')
    plt.grid(axis='y')
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    main()